#include "State.h"
#include "MultibodySystem.h"
#include "RigidBodyNode.h"
#include "FastOps.h"

#include <cstdlib>
#include <iostream>
#include <iomanip>

using namespace std;
using namespace RSIM;

////////////////////////////////////////////////////////////////////////

void State::initialize(MultibodySystem& system){
        #ifdef RSIM_SAFE
        Utils::rsim_assert(nNodes_>0,
                        "Cannot create state for an empty MultibodySystem\n");
        #endif
        BodyNQ_ = new int [nNodes_];
        BodyNU_ = new int [nNodes_];
        nQ_ = 0;
        nU_ = 0;
        
        // Calculates the number of Q's and U's in the system
        for(int i=0;i<nNodes_;++i){
                const RigidBodyNode *Nd = system.getBody(BodyIndex(i+1));
                Nd->data_->setNQ(BodyNQ_+i);
                
                QIndexForNode_.addNode(nQ_);
                nQ_ += BodyNQ_[i];
                
                Nd->data_->setNU(BodyNU_+i);
                UIndexForNode_.addNode(nU_);
                nU_ += BodyNU_[i];
        }
        
        // Initialize internal iterator
        QIndexForNode_.finalize();
        UIndexForNode_.finalize();
                
        // Allocate space for Q,U,UDot of the system
        if(nQ_){
                Q_ = new double [nQ_];
                QDot_ = new double [nQ_];
                U_ = new double [nU_];
                UDot_ = new double [nU_];

                int countQ = 0;
                int countU = 0;
                for(int i=0;i<nNodes_;++i){
                        
                        RigidBody *Body = system.wgetBody(BodyIndex(i+1))->data_;
                        
                        Body->setDefaultQ(Q_+countQ);
                        Body->setDefaultU(U_+countU);
                        
                        {
                                PrescribedMotion *tmp_var = Body->getJoint()->getPrescribedMotion();
                                PM_.addNode(tmp_var);
                                if(tmp_var)
                                        PrescribedMotionIndex_.addNode(i);
                        }

                        if(Body->areQuaternionsInUse()){			

                                // time derivatives of quaternions are computed during the sweeps of the system
                                for(int j=0;j<4;QDot_[countQ + j++]=0.0);
                                
                                // FOr this body, nQ = 4 + nU and nU > 3
                                if(BodyNQ_[i]>4){
                                        for(int j=0;j<(BodyNU_[i]-3);++j){
                                                QDot_[countQ +4+j]=U_[countU + 3 + j];
                                        }
                                }
                        }
                        else{					

                                for(int j=0;j<BodyNU_[i];++j){
                                        QDot_[countQ + j]=U_[countU + j];
                                }
                        }

                        countQ += BodyNQ_[i];
                        countU += BodyNU_[i];
                }
                PM_.finalize();
                QuaternionIndex_.finalize();
                PrescribedMotionIndex_.finalize();
        } // if(nQ_)	
} // void initialize()

////////////////////////////////////////////////////////////////////////

void State::updQ(const double *q,const int& NodeIndex){
        #ifdef RSIM_SAFE
                Utils::rsim_assert(NodeIndex>=0,
                                "State::updQ(), NodeIndex must be greater than or equal to zero\n");
                Utils::rsim_assert(NodeIndex<nNodes_,
                                "State::updQ(), NodeIndex is too high, must be less than the number of nodes in the multibody system\n");
        #endif
        
        const int& id = QIndexForNode_(NodeIndex);
        if(PM_(NodeIndex)){
                #ifdef RSIM_SAFE
                if(q){
                        cout<<"State::updQ(), q ignored since motion is prescribed on this node\n";
                        PM_(NodeIndex)->setQ(*SysData_,Q_+id);
                }
                #endif
        }
        else{
                for(int i=0;i<BodyNQ_[NodeIndex];++i)
                        Q_[id+i] = q[i];
        }
}

////////////////////////////////////////////////////////////////////////

void State::updQ(const double *q, const BodyIndex& BI){
        #ifdef RSIM_SAFE
                assert(BI.id()>=0);
                assert(BI.id()<nNodes_);
        #endif
        this->updQ(q,BI.id());
}

////////////////////////////////////////////////////////////////////////

void State::updStateFromPrescribedMotion(){                
        for(int i=0;i<PrescribedMotionIndex_.getLength();++i){
                int& index = PrescribedMotionIndex_[i];
                PM_(index)->setQ(*SysData_, Q_+QIndexForNode_(index));
                PM_(index)->setU(*SysData_, U_+UIndexForNode_(index));
        }
}

////////////////////////////////////////////////////////////////////////

void State::initTime(TIME *time){
        time_ = time;                
}

////////////////////////////////////////////////////////////////////////

double State::getTime()const{
        return time_->getTime();
}

////////////////////////////////////////////////////////////////////////
                
void State::updU(const double *u,const int& NodeIndex){
        #ifdef RSIM_SAFE
                assert(NodeIndex>=0);
                assert(NodeIndex<nNodes_);
        #endif
        const int& id = UIndexForNode_(NodeIndex);
        if(PM_(NodeIndex)){
                #ifdef RSIM_SAFE
                if(u){
                        cout<<"State::updQ(), q ignored since motion is prescribed on this node\n";
                        PM_(NodeIndex)->setU(*SysData_, U_+id);
                }
                #endif
        }
        else{
                for(int i=0;i<BodyNU_[NodeIndex];++i)
                        U_[id+i] = u[i];
        }
}

////////////////////////////////////////////////////////////////////////

void State::normalizeQuaternions(){	
        for(int i=0;i<QuaternionIndex_.getLength();++i)
                FastOps::normalizeVector(Q_+QIndexForNode_(QuaternionIndex_[i]),4);
}


////////////////////////////////////////////////////////////////////////

State::~State(){
        delete [] Q_;
        delete [] QDot_;
        delete [] U_;
        delete [] UDot_;
        delete [] BodyNQ_;
        delete [] BodyNU_;	
}

////////////////////////////////////////////////////////////////////////

State::State(MultibodySystem& system):	SysData_(NULL),time_(NULL),Q_(NULL),QDot_(NULL),U_(NULL),UDot_(NULL),nQ_(0),nU_(0),
                                        BodyNQ_(NULL),BodyNU_(NULL),nNodes_(system.getNNodes()),
                                        QIndexForNode_(),UIndexForNode_(){
                this->initialize(system);
}

////////////////////////////////////////////////////////////////////////

void State::printQ(){
        cout<<setprecision(16)<<endl;
        for(int i=0;i<nQ_;++i)
                cout<<Q_[i]<<"\t";
        cout<<endl;
}

void State::printQDot(){
        cout<<endl;
        for(int i=0;i<nQ_;++i)
                cout<<QDot_[i]<<"\t";
        cout<<endl;
}

void State::printU(){
        cout<<endl;
        for(int i=0;i<nU_;++i)
                cout<<U_[i]<<"\t";
        cout<<endl;
}

void State::printUDot(){
        cout<<endl;
        for(int i=0;i<nU_;++i)
                cout<<UDot_[i]<<"\t";
        cout<<endl;
}

void State::printUDot1(){
        for(int i=0;i<nU_;++i)
                cout<<UDot_[i]<<"\n";

}

////////////////////////////////////////////////////////////////////////